package com.cb2.algorithm.leetcode;

import java.util.Deque;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;

/**
 * <a href="https://leetcode.cn/problems/kth-smallest-element-in-a-bst/"> 二叉搜索树中第K小的元素(Kth Smallest Element in a BST)</a>
 * <p>给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。</p>
 * <p>
 * <b>示例：</b>
 * <pre>
 * 示例 1：
 *      输入：root = [3,1,4,null,2], k = 1
 *      输出：1
 *
 * 示例 2：
 *      输入：root = [5,3,6,2,4,null,null,1], k = 3
 *      输出：3
 * </pre>
 * </p>
 * <p>
 * <b>提示：</b>
 * <ul>
 *     <li>树中的节点数为 n 。</li>
 *     <li>1 <= k <= n <= 10^4</li>
 *     <li>0 <= Node.val <= 10^4</li>
 * </ul>
 * </p>
 *
 * <b>进阶：如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化算法？</b>
 * ==> 将二叉搜索树转换为平衡二叉搜索树，可以参考<a href="将二叉搜索树转换为平衡二叉搜索树，可以参考「1382. 将二叉搜索树变平衡的官方题解」">「1382. 将二叉搜索树变平衡的官方题解」</a>
 *
 * @author c2b
 * @since 2023/10/19 9:43
 */
public class LC0230KthSmallestElementInBST_M {
    static class Solution {

        /*
         * 方法1适合单次查找
         * 方法2适合多次查找，即进阶中所说的需要频繁地查找第k小的值
         */

        public int kthSmallest(TreeNode root, int k) {
            Deque<TreeNode> stack = new LinkedList<>();
            while (root != null || !stack.isEmpty()) {
                while (root != null) {
                    stack.addFirst(root);
                    root = root.left;
                }
                TreeNode currNode = stack.removeFirst();
                if (--k == 0) {
                    return currNode.val;
                }
                root = currNode.right;
            }
            return -1;
        }

        /**
         * 使用{@literal   Map<TreeNode,int>}结构维护以节点node为根节点的树，节点的个数(含本身)。思路：
         * <p>假设node节点为根节点：进当前节点进行查找：
         * <li>如果node节点左子树的节点数量left = k-1，则第k小的元素即为node。</li>
         * <li>如果node节点左子树的节点数量left < k-1，则第k小的元素一定在node的右子树中，以node的右子节点为根节点查询k = k-left-1。</li>
         * <li>如果node节点左子树的节点数量left > k-1，则第k小的元素一定在node的左子树中，以node的左子节点为根节点查询k。</li>
         * </p>
         */
        public int kthSmallest2(TreeNode root, int k) {
            MyBST bst = new MyBST(root);
            return bst.kthSmallest(k);
        }

        private static class MyBST {
            TreeNode root;
            /**
             * 以key为根的树中节点数量映射
             */
            Map<TreeNode, Integer> tree2nodeNumMap;

            public MyBST(TreeNode root) {
                this.root = root;

                this.tree2nodeNumMap = new HashMap<>();
                calculateNodeNum(root);
            }

            private int calculateNodeNum(TreeNode node) {
                if (node == null) {
                    return 0;
                }
                tree2nodeNumMap.put(node, calculateNodeNum(node.left) + calculateNodeNum(node.right) + 1);
                return tree2nodeNumMap.get(node);
            }

            private int kthSmallest(int k) {
                TreeNode node = root;
                while (node != null) {
                    // 获取其左子树的节点数量
                    int left = tree2nodeNumMap.getOrDefault(node.left, 0);
                    // 如果左子树的节点数量left == k-1，则node节点即为第k小的元素
                    if (left == k - 1) {
                        return node.val;
                    }
                    // 如果左子树的节点数量left < k-1，则第k小的元素在node的右子树中，在其右子树中查询第k-left-1小的元素
                    if (left < k - 1) {
                        node = node.right;
                        k = k - left - 1;
                    }
                    // 如果左子树的节点数量left > k-1，则第k小的元素在node的左子树中，在其左子树中继续查询第k小的元素
                    else {
                        node = node.left;
                    }
                }
                return -1;
            }
        }
    }

    public static void main(String[] args) {
        Solution solution = new Solution();

        TreeNode node1 = new TreeNode(3);
        node1.left = new TreeNode(1);
        node1.left.right = new TreeNode(2);
        node1.right = new TreeNode(4);
        System.out.println(solution.kthSmallest(node1, 1)); // 1
        System.out.println(solution.kthSmallest2(node1, 1)); // 1

        TreeNode node2 = new TreeNode(5);
        node2.left = new TreeNode(3);
        node2.right = new TreeNode(6);
        node2.left.left = new TreeNode(2);
        node2.left.left.left = new TreeNode(1);
        node2.left.right = new TreeNode(4);
        System.out.println(solution.kthSmallest(node2, 3)); // 3
        System.out.println(solution.kthSmallest2(node2, 3)); // 3
    }
}